#!/usr/bin/python3
# -*- coding: utf-8 -*-

import sys
# sys.path.append('/opt/work/caffe/python')
sys.path.insert(0, '.')

import random
import argparse
import os
import numpy as np
from collections import OrderedDict
import struct
import cv2
from math import cos, sin

# SLEEP_TIME = 40
SLEEP_TIME = 1

base_path = r'/rootfs/media/kasim/DataSet/Record'
out_path = r'/rootfs/media/kasim/DataSet/Record'
video_file_list = r'/rootfs/media/kasim/DataSet/Record/face_video_list_bbox_count.attr.txt'

# base_path = r'/rootfs/media/kasim/DataSet/RecordEx'
# out_path = r'/rootfs/media/kasim/DataSet/RecordEx'
# video_file_list = r'/rootfs/media/kasim/DataSet/RecordEx/Face_video_list_bbox_count.attr.txt'

SHOW = True
DISPLAY_WIDTH = 1280
DISPLAY_HEIGHT = 720

min_good_frame_count = 10

# VERY_GOOD = 0
# GOOD = 1
# BAD = 2
# VERY_BAD = 3
filter_quality = 3
filter_quality_thr = 0.25

min_pitch = -30
max_pitch = 30
min_roll = -45
max_roll = 45
min_yaw = -30
max_yaw = 30

min_angle_score = 0.2
max_angle_diff = 5
max_face_count = 1

# def draw_frame(image, frame_num, total_frames_num):
#
#     text_scale = max(1.0, image.shape[1] / 800.)
#
#     cv2.putText(image, 'frame: %d/%d' % (frame_num, total_frames_num),
#                 (0, int(15 * text_scale)), cv2.FONT_HERSHEY_PLAIN, text_scale, (0, 0, 255), thickness=2)


def draw_axis(image, yaw, pitch, roll, tdx, tdy, size=100):
    tdx = int(round(tdx))
    tdy = int(round(tdy))

    LINE_THICKNESS = 4
    # Convert from angular to radian degree
    pitch = pitch * np.pi / 180
    yaw = -(yaw * np.pi / 180)
    roll = roll * np.pi / 180

    # X-Axis pointing to right. drawn in red
    x1 = size * (cos(yaw) * cos(roll)) + tdx
    y1 = size * (cos(pitch) * sin(roll) + cos(roll) * sin(pitch) * sin(yaw)) + tdy

    # Y-Axis | drawn in green
    x2 = size * (-cos(yaw) * sin(roll)) + tdx
    y2 = size * (cos(pitch) * cos(roll) - sin(pitch) * sin(yaw) * sin(roll)) + tdy

    # Z-Axis (out of the screen) drawn in blue
    x3 = size * (sin(yaw)) + tdx
    y3 = size * (-cos(yaw) * sin(pitch)) + tdy

    x1 = int(round(x1))
    y1 = int(round(y1))
    x2 = int(round(x2))
    y2 = int(round(y2))
    x3 = int(round(x3))
    y3 = int(round(y3))

    cv2.line(image, (tdx, tdy), (x1, y1), (0, 0, 255), thickness=LINE_THICKNESS)
    cv2.line(image, (tdx, tdy), (x2, y2), (0, 255, 0), thickness=LINE_THICKNESS)
    cv2.line(image, (tdx, tdy), (x3, y3), (255, 0, 0), thickness=LINE_THICKNESS)


def visualize_restuls(image, attributes, frame_num, total_frames_num):
    text_scale = max(1.0, image.shape[1] / 800.)

    cv2.putText(image, 'frame: %d/%d' % (frame_num, total_frames_num),
                (0, int(15 * text_scale)), cv2.FONT_HERSHEY_PLAIN, text_scale, (0, 0, 255), thickness=2)

    if attributes is not None:
        for attribute in attributes:
            pitch = attribute['pitch'][0]
            yaw = attribute['yaw'][0]
            roll = attribute['roll'][0]

            xmin, ymin, xmax, ymax = attribute['bboxes']

            draw_axis(image, yaw, pitch, roll, tdx=(xmin + xmax) / 2, tdy=(ymin + ymax) / 2, size=(ymax - ymin) / 2)

            xmin, ymin, xmax, ymax = list(map(int, map(round, attribute['bboxes'])))
            cv2.rectangle(image, (xmin, ymin), (xmax, ymax), (0, 255, 0), thickness=2)

    return image


def main():
    if not os.path.exists(out_path):
        os.makedirs(out_path)
    os.system('chmod a+wr {}'.format(out_path))

    video_list_dict = OrderedDict()
    with open(video_file_list, 'r') as file:
        for line in file.readlines():
            line_info = line.strip().split(',')
            file_path = line_info[0]
            face_count = int(line_info[1])
            video_list_dict[file_path] = face_count

    _base_path = os.path.splitext(video_file_list)[0]
    title_file_path = _base_path + '.title.txt'
    with open(title_file_path, 'r') as file:
        title_str = file.readline()
        attr_item_count_str = file.readline()
        titles = title_str.strip().split(';')
        attr_item_counts = attr_item_count_str.strip().split(';')
        attr_item_counts = list(map(int, attr_item_counts))
        total_attr_item_count = sum(attr_item_counts)
        attr_count = len(titles)

    attr_item_offsets = [1]
    offset = 1
    for attr_item_count in attr_item_counts:
        offset += attr_item_count
        attr_item_offsets.append(offset)
    attrs_size = 4 + total_attr_item_count * 4
    video_attributes_dict = {}
    data_file_path = _base_path + '.dat'
    with open(data_file_path, 'rb') as file:
        for video_name, face_count in video_list_dict.items():
            attributes_dict = {}
            for _ in range(face_count):
                dat = file.read(attrs_size)
                datas = struct.unpack('1i{}f'.format(total_attr_item_count), dat)
                frame_id = datas[0]
                attribute = {}
                for i, title in enumerate(titles):
                    start_offset = attr_item_offsets[i]
                    end_offset = attr_item_offsets[i+1]
                    attribute[title] = datas[start_offset:end_offset]
                attributes = attributes_dict.get(frame_id, None)
                if attributes is None:
                    attributes_dict[frame_id] = [attribute]
                else:
                    attributes.append(attribute)
            video_attributes_dict[video_name] = attributes_dict
            # break

    if SHOW:
        for video_name, attributes_dict in video_attributes_dict.items():
            full_video_path = os.path.join(base_path, video_name)
            cap = cv2.VideoCapture(full_video_path)
            total_frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
            imWidth = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
            imHeight = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
            frameRate = int(cap.get(cv2.CAP_PROP_FPS))
            frame_count = 0
            while True:
                grabbed, img = cap.read()
                if not grabbed:
                    break

                attributes = attributes_dict.get(frame_count, None)
                frame_count += 1
                visualize_restuls(img, attributes, frame_count, total_frame_count)
                if (img.shape[0] != DISPLAY_HEIGHT) or (img.shape[1] != DISPLAY_WIDTH):
                    img = cv2.resize(img, (DISPLAY_WIDTH, DISPLAY_HEIGHT))
                cv2.imshow('image', img)
                if attributes is None:
                    key = cv2.waitKey(1)
                else:
                    key = cv2.waitKey()
                if key == 27 or key == ord('q' or key == ord('Q')):
                    return
    else:
        out_file_path = os.path.join(out_path, 'video_attr_filter_file_list.txt')
        out_file = open(out_file_path, 'w')
        total_file_count = len(video_attributes_dict)
        file_count = 0
        select_file_count = 0
        for video_name, attributes_dict in video_attributes_dict.items():
            skip_video = False
            good_frame_count = 0
            file_count += 1
            if file_count % 100 == 0:
                print('Process File Count: {}/{}, Select File Count: {}'.format(file_count, total_file_count, select_file_count))
            for frame_id, attributes in attributes_dict.items():
                if len(attributes) > max_face_count:
                    skip_video = True
                    break
                for attribute in attributes:
                    quality = attribute['quality']
                    quality = np.array(quality, dtype=np.float32)
                    max_idx = np.argmax(quality)
                    if (max_idx == filter_quality) and (quality[max_idx] >= filter_quality_thr):
                        continue

                    pitch = attribute['pitch'][0]
                    yaw = attribute['yaw'][0]
                    roll = attribute['roll'][0]

                    if pitch < min_pitch or max_pitch < pitch or roll < min_roll or max_roll < roll or yaw < min_yaw or max_yaw < yaw:
                        continue

                    bin_pitch = attribute['bin_pitch']
                    bin_yaw = attribute['bin_yaw']
                    bin_roll = attribute['bin_roll']

                    bin_pitch_score = float(bin_pitch[1])
                    bin_pitch = float(bin_pitch[0])
                    bin_roll_score = float(bin_roll[1])
                    bin_roll = float(bin_roll[0])
                    bin_yaw_score = float(bin_yaw[1])
                    bin_yaw = float(bin_yaw[0])

                    pitch_diff = abs(bin_pitch - pitch)
                    roll_diff = abs(bin_roll - roll)
                    yaw_diff = abs(bin_yaw - yaw)
                    if max_angle_diff < pitch_diff or max_angle_diff < roll_diff or max_angle_diff < yaw_diff:
                         continue

                    if bin_pitch_score <= min_angle_score or bin_roll_score <= min_angle_score or bin_yaw_score <= min_angle_score:
                        continue

                    good_frame_count += 1
                    break
            if skip_video:
                continue
            frame_count = len(attributes_dict)
            if good_frame_count > min_good_frame_count:
                out_info_str = '{} {} {}\n'.format(video_name, frame_count, good_frame_count)
                out_file.write(out_info_str)
                select_file_count += 1
        print('Process File Count: {}/{}, Select File Count: {}'.format(file_count, total_file_count, select_file_count))
        out_file.close()
        os.system('chmod a+wr {}'.format(out_file_path))

    print('Finish!')


if __name__ == '__main__':
    main()
